class Solution(object):
    def findDiagonalOrder(self, mat):
        """
        :type mat: List[List[int]]
        :rtype: List[int]
        """
        m = len(mat)
        n = len(mat[0])
        state = 1
        index = 0
        out = [0] * m * n
        for k in range(m + n - 1):
            if state:
                if k < m:
                    i = k
                else:
                    i = m - 1
                while i >= 0 and k - i < n:
                    out[index] = mat[i][k - i]
                    index += 1
                    i -= 1
                state = 0
            else:
                if k < n:
                    j = k
                else:
                    j = n - 1
                while j >= 0 and k - j < m:
                    out[index] = mat[k - j][j]
                    index += 1
                    j -= 1
                state = 1
        return out


data = Solution()
mat = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
print(data.findDiagonalOrder(mat))
